#include "Multigrid.h"

#include <math.h>

#include "basis.h"
#include "Iteration.h"

#include "DimensionSelect.h"

#include "DebugHelper.h"

using namespace Eigen;
using namespace std;

#define EXPAND(...) __VA_ARGS__
#define EXPAND_POINT(point) EXPAND((point).first, (point).second)

runningLog::runningLog(double err, int count) : err(err), count(count)
{
}
runningLog::runningLog() : runningLog(0, 0) {}
runningLog::runningLog(const Eigen::VectorXd &err, int count) : runningLog::runningLog(vectorXdInftyNum(err), count) {}

void Multigrid::updateDim()
{
#define CASE_DIM(dim)                       \
    case dim:                               \
        restriction = GET_Restriction(dim); \
        prolongate = GET_Prolongate(dim);   \
        break;

    switch (dim)
    {
        CASE_DIM(One);
        CASE_DIM(Two);
    default:
        break;
    }
#undef CASE_DIM
}

int Multigrid::size() { return (int)pow((1 << (degree + 1)) + 1, dim); }

Multigrid::Multigrid(int degree, Function f, Dimension d) : degree(degree), f(f), dim(d), cycletype(Vcycle), rtype(Operators::Injection), itype(Operators::Linear), btype(Dirichlet) { updateDim(); }

void Multigrid::setDimension(Dimension d)
{
    dim = d;
    updateDim();
}

void Multigrid::setCycleType(CycleChoice c) { cycletype = c; }
void Multigrid::setRestrictionType(Operators::Restriction s) { rtype = s; }
void Multigrid::setInterpolationType(Operators::Interpolation t) { itype = t; }
void Multigrid::setBoundtype(BoundType b) { btype = b; }

#define SET_AH_FUN(dim) setAh_##dim

inline void SET_AH_FUN(One)(int size, _In_ Function &f, BoundType btype, Dimension dim, SparseMatrix<double> &Ah, Eigen::VectorXd *pValues = nullptr)
{
    int n = getN(size, dim);
    double h = 1.0 / n;
    switch (btype)
    {
    case BoundType::Dirichlet:
        for (int i = 0; i < size; i++)
        {
            if (i == 0 || i == size - 1)
            {
                Ah.coeffRef(i, i) = 1;
                if (pValues)
                    (*pValues)(i) = f(i * h);
            }
            else
            {
                Ah.coeffRef(i, i - 1) = -1. / (h * h);
                Ah.coeffRef(i, i + 1) = -1. / (h * h);
                Ah.coeffRef(i, i) = 2. / (h * h);
                if (pValues)
                    (*pValues)(i) = f.laplace(i * h);
            }
        }
        break;
    case BoundType::Neumann:
        for (int i = 0; i < size; i++)
        {
            if (i == 0)
            {
                Ah.coeffRef(i, i) = 1;
                if (pValues)
                    (*pValues)(i) = f(0);
            }
            else if (i == size - 1)
            {
                Ah.coeffRef(i, i) = 1.0 / h;
                Ah.coeffRef(i, i - 1) = -1.0 / h;
                if (pValues)
                    (*pValues)(i) = f.dev1d(1.0);
            }
            else
            {
                Ah.coeffRef(i, i) = 2.0 / (h * h);
                Ah.coeffRef(i, i - 1) = -1.0 / (h * h);
                Ah.coeffRef(i, i + 1) = -1.0 / (h * h);
                if (pValues)
                    (*pValues)(i) = f.laplace(i * h);
            }
        }
    default:
        break;
    }
}

inline void SET_AH_FUN(Two)(int size, _In_ Function &f, BoundType btype, Dimension dim, SparseMatrix<double> &Ah, Eigen::VectorXd *pValues = nullptr)
{
    int n = getN(size, dim);
    double h = 1.0 / n;
    vector<vector<pair<double, double>>> points; // 二维都用n存，因为点的个数是（n+1）^2
    for (int i = 0; i <= n; i++)
    {
        vector<pair<double, double>> tmp;
        for (int j = 0; j <= n; j++)
            tmp.push_back(make_pair(j * h, i * h));
        points.push_back(tmp);
    }
    int cnt = 0;
    switch (btype)
    {
    case BoundType::Dirichlet:
        for (int i = 0; i <= n; i++)
        {
            for (int j = 0; j <= n; j++)
            {
                if (isinsquare(points[i][j]) != Insqaure::Inside)
                {
                    Ah.coeffRef(cnt, cnt) = 1;
                    if (pValues)
                        (*pValues)(cnt) = f(EXPAND_POINT(points[i][j]));
                }
                else
                {
                    Ah.coeffRef(cnt, cnt - 1) = -1. / (h * h);
                    Ah.coeffRef(cnt, cnt + 1) = -1. / (h * h);
                    Ah.coeffRef(cnt, cnt - n - 1) = -1. / (h * h);
                    Ah.coeffRef(cnt, cnt + n + 1) = -1. / (h * h);
                    Ah.coeffRef(cnt, cnt) = 4. / (h * h);
                    if (pValues)
                        (*pValues)(cnt) = f.laplace(EXPAND_POINT(points[i][j]));
                }
                cnt++;
            }
        }
        break;
    case BoundType::Neumann:
        for (int i = 0; i <= n; i++)
        {
            for (int j = 0; j <= n; j++)
            {
                double *devxy = f.dev(points[i][j].first, points[i][j].second);
                double devx = devxy[0];
                double devy = devxy[1];
                delete[] devxy;
                if (isinsquare(points[i][j]) == Insqaure::Vertex)
                {
                    Ah.coeffRef(cnt, cnt) = 1;
                    if (pValues)
                        (*pValues)(cnt) = f(EXPAND_POINT(points[i][j]));
                }
                else if (isinsquare(points[i][j]) == Insqaure::Inside)
                {
                    Ah.coeffRef(cnt, cnt) = 4;
                    Ah.coeffRef(cnt, cnt + 1) = -1;
                    Ah.coeffRef(cnt, cnt - 1) = -1;
                    Ah.coeffRef(cnt, cnt + n + 1) = -1;
                    Ah.coeffRef(cnt, cnt - n - 1) = -1;
                    if (pValues)
                        (*pValues)(cnt) = f.laplace(EXPAND_POINT(points[i][j])) * h * h;
                }
                else if (isinsquare(points[i][j]) == Insqaure::Left)
                {
                    Ah.coeffRef(cnt, cnt) = 4;
                    Ah.coeffRef(cnt, cnt + 1) = -2;
                    Ah.coeffRef(cnt, cnt + n + 1) = -1;
                    Ah.coeffRef(cnt, cnt - n - 1) = -1;
                    if (pValues)
                        (*pValues)(cnt) = h * h * f(EXPAND_POINT(points[i][j])) - devx * 2 * h;
                }
                else if (isinsquare(points[i][j]) == Insqaure::Right)
                {
                    Ah.coeffRef(cnt, cnt) = 4;
                    Ah.coeffRef(cnt, cnt - 1) = -2;
                    Ah.coeffRef(cnt, cnt + n + 1) = -1;
                    Ah.coeffRef(cnt, cnt - n - 1) = -1;
                    if (pValues)
                        (*pValues)(cnt) = h * h * f(EXPAND_POINT(points[i][j])) - devx * 2 * h;
                }
                else if (isinsquare(points[i][j]) == Insqaure::Up)
                {
                    Ah.coeffRef(cnt, cnt) = 4;
                    Ah.coeffRef(cnt, cnt - 1) = -1;
                    Ah.coeffRef(cnt, cnt + 1) = -1;
                    Ah.coeffRef(cnt, cnt - n - 1) = -2;
                    if (pValues)
                        (*pValues)(cnt) = h * h * f(EXPAND_POINT(points[i][j])) - devy * 2 * h;
                }
                else if (isinsquare(points[i][j]) == Insqaure::Down)
                {
                    Ah.coeffRef(cnt, cnt) = 4;
                    Ah.coeffRef(cnt, cnt + 1) = -1;
                    Ah.coeffRef(cnt, cnt - 1) = -1;
                    Ah.coeffRef(cnt, cnt + n + 1) = -2;
                    if (pValues)
                        (*pValues)(cnt) = h * h * f(EXPAND_POINT(points[i][j])) - devy * 2 * h;
                }
                cnt++;
            }
        }
        break;
    default:
        break;
    }
}

SparseMatrix<double> Multigrid::setgrid(int size, Eigen::VectorXd *pValues)
{
    if (pValues)
        *pValues = VectorXd::Zero(size);
    SparseMatrix<double> Ah(size, size);

#define CASE_DIM(dim)                       \
    case dim:                               \
        SET_AH_FUN(dim)                     \
        (size, f, btype, dim, Ah, pValues); \
        break;

    switch (dim)
    {
        CASE_DIM(One);
        CASE_DIM(Two);
        break;
    default:
        break;
    }
    return Ah;
#undef CASE_DIM
}

VectorXd Multigrid::V_cycle(const VectorXd &_v, const VectorXd &f, int n1, int n2, double w)
{
    SparseMatrix<double> Ah = setgrid(_v.size());
    VectorXd v = relaxiteration(Ah, f, _v, n1, w);
    int n = getN(v.size(), dim);
    if (n > 2)
    {
        int newsize = (n >> 1) + 1;
        VectorXd v2 = VectorXd::Zero((int)pow(newsize, dim));
        VectorXd r = f - Ah * v;
        VectorXd tmp = restriction(rtype, r);
        v2 = V_cycle(v2, tmp, n1, n2, w);
        v = prolongate(itype, v2) + v;
    }
    v = relaxiteration(Ah, f, v, n2, w);
    return v;
}

VectorXd Multigrid::FMG(const VectorXd &f, int n1, int n2, double w)
{
    VectorXd vtmp;
    int n = getN(f.size(), dim);
    if (n != 2)
    {
        VectorXd ftmp = restriction(rtype, f);
        vtmp = FMG(ftmp, n1, n2, w);
        vtmp = prolongate(itype, vtmp);
    }
    else
    {
        vtmp = VectorXd::Zero(f.size());
    }
    return V_cycle(vtmp, f, n1, n2, w);
}

void analytical(const SparseMatrix<double> &A, const Eigen::VectorXd &f, _Out_ Eigen::VectorXd &solution)
{
    SparseLU<SparseMatrix<double>> solver;
    solver.compute(A);
    solution = solver.solve(f);
}

std::tuple<Eigen::VectorXd, Eigen::VectorXd> Multigrid::Possion_BVP_Multigrid(int nu1, int nu2,
                                                                              int max_iter, double epsi, double w, _Out_ runningLog *pLog)
{
    Eigen::VectorXd f;
    SparseMatrix<double> A = setgrid(size(), &f);
    VectorXd res = VectorXd::Zero(f.size());
    VectorXd solution = VectorXd::Zero(f.size());
    VectorXd err;
    runningLog l;
#define LOOP(GET_RES, DONE)                      \
    for (l.count = 0; (l.count) < max_iter; l++) \
    {                                            \
        res = GET_RES;                           \
        analytical(A, f, solution);              \
        err = res - solution;                    \
        l.updateError(err);                      \
        if (err <= epsi)                         \
            break;                               \
        DONE;                                    \
    }
    switch (cycletype)
    {
    case CycleChoice::Vcycle:
        LOOP(V_cycle(res, f, nu1, nu2, w), );
        break;
    case CycleChoice::FMG:
        LOOP(FMG(f, nu1, nu2, w), f = res);
        break;
    default:
        break;
    }
#undef LOOP
    l.count = min(l.count + 1, max_iter);
    std::cout << "err = " << l.err << (l.err <= epsi ? " < epsi: " : " by ") << l.count << " iteration" << endl;
    if (pLog)
        *pLog = l;
    return std::make_tuple(res, solution);
}

bool operator==(std::pair<double, double> p1, std::pair<double, double> p2)
{
    if (fabs(p1.first - p2.first) > 1e-12 || fabs(p1.first - p2.first) > 1e-12)
    {
        return false;
    }
    return true;
}

Insqaure isinsquare(double x, double y)
{
    if (x > 1e-12 && x < 1 - 1e-12 && y > 1e-12 && y < 1 - 1e-12)
        return Insqaure::Inside;
    else if (x < 1e-12 && y > 1e-12 && y < 1 - 1e-12) // left
        return Insqaure::Left;
    else if (abs(1.0 - y) < 1e-12 && x > 1e-12 && x < 1 - 1e-12) // up
        return Insqaure::Right;
    else if (abs(1.0 - x) < 1e-12 && y > 1e-12 && y < 1 - 1e-12) // right
        return Insqaure::Up;
    else if (y < 1e-12 && x > 1e-12 && x < 1 - 1e-12) // bottom
        return Insqaure::Down;
    else
        return Insqaure::Vertex;
}

inline Insqaure isinsquare(std::pair<double, double> p)
{
    return isinsquare(p.first, p.second);
}

inline double distance(std::pair<double, double> p1, std::pair<double, double> p2)
{
    return sqrt((p1.first - p2.first) * (p1.first - p2.first) + (p1.second - p2.second) * (p1.second - p2.second));
}